#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
继电器MQTT测试程序
用于测试ESP32继电器控制板的MQTT协议功能
协议格式：[0x55, 地址, 功能码, 数据1, 数据2, 数据3, 数据4, 校验和]
"""

import paho.mqtt.client as mqtt
import time
import sys
import struct

class RelayMqttTester:
    def __init__(self, broker_host="110.41.147.192", broker_port=1883, device_mac="3c8a1fd51e48"):
        """
        初始化MQTT测试客户端
        
        Args:
            broker_host: MQTT服务器地址
            broker_port: MQTT服务器端口
            device_mac: 设备MAC地址
        """
        self.broker_host = broker_host
        self.broker_port = broker_port
        self.device_mac = device_mac
        
        # MQTT主题
        self.publish_topic = f"/c/{device_mac}/get"  # 发送给设备
        self.subscribe_topic = f"/c/{device_mac}/send"  # 接收设备回复
        
        # 初始化MQTT客户端
        self.client = mqtt.Client()
        self.client.on_connect = self.on_connect
        self.client.on_message = self.on_message
        self.client.on_disconnect = self.on_disconnect
        
        print(f"🔧 继电器MQTT测试器初始化")
        print(f"📡 MQTT服务器: {broker_host}:{broker_port}")
        print(f"🎯 设备MAC: {device_mac}")
        print(f"📤 发布主题: {self.publish_topic}")
        print(f"📥 订阅主题: {self.subscribe_topic}")
        print("-" * 60)

    def calculate_checksum(self, data):
        """计算校验和"""
        return sum(data) & 0xFF

    def create_protocol_packet(self, address=0x01, function_code=0x10, data1=0x00, data2=0x00, data3=0x00, data4=0x00):
        """
        创建8字节协议数据包
        
        Args:
            address: 设备地址
            function_code: 功能码
            data1-data4: 数据字节
            
        Returns:
            bytes: 8字节协议数据包
        """
        packet = [0x55, address, function_code, data1, data2, data3, data4, 0x00]
        packet[7] = self.calculate_checksum(packet[:7])
        return bytes(packet)

    def on_connect(self, client, userdata, flags, rc):
        """MQTT连接回调"""
        if rc == 0:
            print("✅ MQTT连接成功!")
            client.subscribe(self.subscribe_topic)
            print(f"📥 已订阅主题: {self.subscribe_topic}")
        else:
            print(f"❌ MQTT连接失败，错误码: {rc}")

    def on_message(self, client, userdata, msg):
        """MQTT消息接收回调"""
        print(f"\n📨 收到设备回复:")
        print(f"   主题: {msg.topic}")
        print(f"   长度: {len(msg.payload)} 字节")
        
        # 解析8字节回复数据
        if len(msg.payload) == 8:
            data = list(msg.payload)
            print(f"   数据: {' '.join([f'0x{b:02X}' for b in data])}")
            
            # 解析协议内容
            header = data[0]
            address = data[1] 
            function_code = data[2]
            relay_state1 = data[5]  # 第16-9路继电器状态
            relay_state2 = data[6]  # 第8-1路继电器状态
            checksum = data[7]
            
            print(f"   协议头: 0x{header:02X}")
            print(f"   地址: 0x{address:02X}")
            print(f"   功能码: 0x{function_code:02X}")
            print(f"   继电器状态1 (16-9路): 0x{relay_state1:02X} ({relay_state1:08b})")
            print(f"   继电器状态2 (8-1路): 0x{relay_state2:02X} ({relay_state2:08b})")
            print(f"   校验和: 0x{checksum:02X}")
            
            # 解析具体继电器状态
            print(f"   🔍 继电器状态详情:")
            for i in range(14,8,-1):
                state = "开启" if (relay_state1 & (1 << i-9)) else "关闭"
                print(f"      继电器{i}: {state}")
            for i in range(8,0,-1):
                state = "开启" if (relay_state2 & (1 << i-1)) else "关闭"
                print(f"      继电器{i}: {state}")
        else:
            print(f"   原始数据: {msg.payload.hex()}")

    def on_disconnect(self, client, userdata, rc):
        """MQTT断开连接回调"""
        print(f"🔌 MQTT连接断开，错误码: {rc}")

    def connect(self):
        """连接MQTT服务器"""
        try:
            print(f"🔄 正在连接MQTT服务器...")
            self.client.connect(self.broker_host, self.broker_port, 60)
            self.client.loop_start()
            time.sleep(2)  # 等待连接建立
            return True
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return False

    def disconnect(self):
        """断开MQTT连接"""
        self.client.loop_stop()
        self.client.disconnect()
        print("👋 已断开MQTT连接")

    def send_command(self, packet, description=""):
        """发送命令到设备"""
        try:
            packet_hex = ' '.join([f'0x{b:02X}' for b in packet])
            print(f"\n📤 发送命令: {description}")
            print(f"   数据: {packet_hex}")
            
            result = self.client.publish(self.publish_topic, packet)
            if result.rc == mqtt.MQTT_ERR_SUCCESS:
                print(f"   ✅ 命令发送成功")
                time.sleep(0.5)  # 等待设备响应
                return True
            else:
                print(f"   ❌ 命令发送失败，错误码: {result.rc}")
                return False
        except Exception as e:
            print(f"   ❌ 发送异常: {e}")
            return False

    # 测试函数集合
    def test_read_status(self):
        """测试读取继电器状态"""
        packet = self.create_protocol_packet(function_code=0x10)
        return self.send_command(packet, "读取继电器状态 (功能码0x10)")

    def test_open_relay(self, relay_num):
        """测试打开指定继电器"""
        if not (1 <= relay_num <= 14):
            print(f"❌ 继电器编号错误: {relay_num} (支持1-14)")
            return False
        packet = self.create_protocol_packet(function_code=0x12, data4=relay_num)
        return self.send_command(packet, f"打开继电器{relay_num} (功能码0x12)")

    def test_close_relay(self, relay_num):
        """测试关闭指定继电器"""
        if not (1 <= relay_num <= 14):
            print(f"❌ 继电器编号错误: {relay_num} (支持1-14)")
            return False
        packet = self.create_protocol_packet(function_code=0x11, data4=relay_num)
        return self.send_command(packet, f"关闭继电器{relay_num} (功能码0x11)")

    def test_open_all(self):
        """测试打开所有继电器"""
        packet = self.create_protocol_packet(function_code=0x13, data4=0xFF)
        return self.send_command(packet, "打开所有继电器 (功能码0x13)")

    def test_close_all(self):
        """测试关闭所有继电器"""
        packet = self.create_protocol_packet(function_code=0x13, data4=0x00)
        return self.send_command(packet, "关闭所有继电器 (功能码0x13)")

    def test_open_delay(self, relay_num, delay_ms):
        """测试打开延时关闭"""
        if not (1 <= relay_num <= 14):
            print(f"❌ 继电器编号错误: {relay_num} (支持1-14)")
            return False
        
        # 将延时时间转换为3字节 (高位在前)
        data1 = (delay_ms >> 16) & 0xFF
        data2 = (delay_ms >> 8) & 0xFF
        data3 = delay_ms & 0xFF
        
        packet = self.create_protocol_packet(function_code=0x21, data1=data1, data2=data2, data3=data3, data4=relay_num)
        return self.send_command(packet, f"继电器{relay_num} 打开延时{delay_ms}ms后关闭 (功能码0x21)")

    def test_close_delay(self, relay_num, delay_ms):
        """测试关闭延时打开"""
        if not (1 <= relay_num <= 14):
            print(f"❌ 继电器编号错误: {relay_num} (支持1-14)")
            return False
        
        # 将延时时间转换为3字节 (高位在前)
        data1 = (delay_ms >> 16) & 0xFF
        data2 = (delay_ms >> 8) & 0xFF
        data3 = delay_ms & 0xFF
        
        packet = self.create_protocol_packet(function_code=0x22, data1=data1, data2=data2, data3=data3, data4=relay_num)
        return self.send_command(packet, f"继电器{relay_num} 关闭延时{delay_ms}ms后打开 (功能码0x22)")

    def test_restart_system(self):
        """测试重启系统"""
        packet = self.create_protocol_packet(function_code=0x81)
        return self.send_command(packet, "重启系统 (功能码0x81)")

def show_menu():
    """显示测试菜单"""
    print("\n" + "="*60)
    print("🎮 继电器MQTT测试菜单")
    print("="*60)
    print("1.  读取继电器状态")
    print("2.  打开指定继电器")
    print("3.  关闭指定继电器")
    print("4.  打开所有继电器")
    print("5.  关闭所有继电器")
    print("6.  继电器打开延时关闭测试")
    print("7.  继电器关闭延时打开测试")
    print("8.  自动化测试所有继电器")
    print("9.  发送自定义命令")
    print("10. 重启系统")
    print("0.  退出程序")
    print("-"*60)

def run_auto_test(tester):
    """运行自动化测试"""
    print("\n🚀 开始自动化测试...")
    
    # 1. 读取初始状态
    print("\n📋 步骤1: 读取初始状态")
    tester.test_read_status()
    time.sleep(2)
    
    # 2. 测试单个继电器控制
    print("\n📋 步骤2: 测试单个继电器控制")
    for relay in [1, 3, 5, 7]:
        print(f"\n🔧 测试继电器{relay}")
        tester.test_open_relay(relay)
        time.sleep(1)
        tester.test_read_status()
        time.sleep(1)
        tester.test_close_relay(relay)
        time.sleep(1)
    
    # 3. 测试全部继电器控制
    print("\n📋 步骤3: 测试全部继电器控制")
    tester.test_open_all()
    time.sleep(2)
    tester.test_read_status()
    time.sleep(2)
    tester.test_close_all()
    time.sleep(2)
    tester.test_read_status()
    
    # 4. 测试延时功能
    print("\n📋 步骤4: 测试延时功能")
    tester.test_open_delay(1, 3000)  # 继电器1打开3秒后关闭
    time.sleep(1)
    tester.test_close_delay(2, 2000)  # 继电器2关闭2秒后打开
    
    print("\n✅ 自动化测试完成!")

def main():
    """主函数"""
    print("🎯 继电器MQTT测试程序")
    print("📝 请输入设备参数 (直接按Enter使用默认值)")
    
    # 获取用户输入
    broker = input("MQTT服务器地址 [110.41.147.192]: ").strip() or "110.41.147.192"
    port_str = input("MQTT端口 [1883]: ").strip() or "1883"
    mac = input("设备MAC地址 [3c8a1fd51e48]: ").strip() or "3c8a1fd51e48"
    
    try:
        port = int(port_str)
    except ValueError:
        print("❌ 端口号无效，使用默认值1883")
        port = 1883
    
    # 创建测试器
    tester = RelayMqttTester(broker, port, mac)
    
    # 连接MQTT服务器
    if not tester.connect():
        print("❌ 无法连接MQTT服务器，程序退出")
        return
    
    try:
        while True:
            show_menu()
            choice = input("请选择操作 (0-10): ").strip()
            
            if choice == "0":
                break
            elif choice == "1":
                tester.test_read_status()
            elif choice == "2":
                relay_num = int(input("请输入继电器编号 (1-14): "))
                tester.test_open_relay(relay_num)
            elif choice == "3":
                relay_num = int(input("请输入继电器编号 (1-14): "))
                tester.test_close_relay(relay_num)
            elif choice == "4":
                tester.test_open_all()
            elif choice == "5":
                tester.test_close_all()
            elif choice == "6":
                relay_num = int(input("请输入继电器编号 (1-14): "))
                delay_ms = int(input("请输入延时时间 (毫秒): "))
                tester.test_open_delay(relay_num, delay_ms)
            elif choice == "7":
                relay_num = int(input("请输入继电器编号 (1-14): "))
                delay_ms = int(input("请输入延时时间 (毫秒): "))
                tester.test_close_delay(relay_num, delay_ms)
            elif choice == "8":
                run_auto_test(tester)
            elif choice == "9":
                print("请输入8字节命令 (十六进制，用空格分隔):")
                hex_input = input("例如: 55 01 10 00 00 00 00 66\n> ")
                try:
                    hex_bytes = [int(x, 16) for x in hex_input.split()]
                    if len(hex_bytes) == 8:
                        packet = bytes(hex_bytes)
                        tester.send_command(packet, "自定义命令")
                    else:
                        print("❌ 命令长度必须为8字节")
                except ValueError:
                    print("❌ 命令格式错误")
            elif choice == "10":
                confirm = input("⚠️  确认要重启系统吗? (y/N): ")
                if confirm.lower() == 'y':
                    tester.test_restart_system()
            else:
                print("❌ 无效选择")
                
    except KeyboardInterrupt:
        print("\n\n⛔ 用户中断程序")
    except Exception as e:
        print(f"\n❌ 程序异常: {e}")
    finally:
        tester.disconnect()
        print("👋 程序结束")

if __name__ == "__main__":
    main() 